import xarray as xr
import numpy as np
from scipy.interpolate import griddata
import netCDF4 as nc

# 读取地表数据（已裁剪）
era5_surface = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/202001000_c.nc')
sp = era5_surface['sp']  # 地表气压 (Pa)
t2m = era5_surface['t2m']
lon_era5 = era5_surface.longitude.values
lat_era5 = era5_surface.latitude.values

# 读取分层比湿数据（需确保与地表数据维度一致）
era5_q = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/specifichumidity_1000hPa_2020010100.nc')
q = era5_q['q']

# 关键修改1：统一维度名称
# 检查比湿数据维度名称
print("比湿数据维度:", q.dims)  # 示例输出：('time', 'level', 'latitude', 'longitude')
q = q.rename({'level': 'pressure_level'}) if 'level' in q.dims else q

# 关键修改2：对齐坐标
# 确保分层数据与地表数据使用相同的经纬度网格
q = q.interp(longitude=lon_era5, latitude=lat_era5)

# # 转换为相同单位
# sp_hPa = sp / 100  # Pa -> hPa
# levels = q.pressure_level.values.astype(float)  # 确保为浮点型
#
#
# # 创建匹配函数
# def find_nearest_level(p):
#     return levels[np.abs(levels - p).argmin()]
#
#
# # 矢量计算最近气压层
# nearest_levels = xr.apply_ufunc(
#     find_nearest_level,
#     sp_hPa,
#     input_core_dims=[[]],
#     output_core_dims=[[]],
#     vectorize=True,
#     dask='parallelized',
#     output_dtypes=[levels.dtype]
# )

# 提取对应层比湿
# surface_q = q.sel(pressure_level=nearest_levels)
surface_q = q.sel(pressure_level=1000, drop=True)

# # 处理极端情况
# min_level = levels.min()
# max_level = levels.max()
# surface_q = surface_q.where(
#     (sp_hPa >= min_level) & (sp_hPa <= max_level),
#     q.sel(pressure_level=min_level)
# )

# 读取FY4A坐标
coord_file = nc.Dataset('/home/liudd/data_preprocessing/FY4A_coordinates.nc', 'r')
lon_fy4a = coord_file.variables['lon'][:, :].T
lat_fy4a = coord_file.variables['lat'][:, :].T
lon_fy4a_360 = lon_fy4a % 360


# 优化插值函数
def interpolate_to_fy4a(data_var):
    # 生成ERA5网格点
    lon_grid, lat_grid = np.meshgrid(lon_era5, lat_era5)
    points = np.column_stack((lon_grid.ravel(), lat_grid.ravel()))
    values = data_var.values.ravel()

    # 准备目标点
    target = np.column_stack((lon_fy4a_360.ravel(), lat_fy4a.ravel()))
    valid = ~np.isnan(target).any(axis=1)

    # 执行插值
    result = np.full(lon_fy4a.shape, np.nan)
    result.ravel()[valid] = griddata(
        points,
        values,
        target[valid],
        method='linear',
        fill_value=np.nan
    )
    return result


# 分步插值验证
try:
    result_t2m = interpolate_to_fy4a(t2m)
    result_sp = interpolate_to_fy4a(sp)
    result_q = interpolate_to_fy4a(surface_q)
except Exception as e:
    print(f"插值错误: {str(e)}")
    # 维度验证
    print(f"T2M形状: {t2m.shape}")
    print(f"SP形状: {sp.shape}")
    print(f"Q形状: {surface_q.shape}")
    exit()

# 保存结果
ds = xr.Dataset(
    {
        "t2m": (("y", "x"), result_t2m),
        "sp": (("y", "x"), result_sp),
        "q_surface": (("y", "x"), result_q)
    },
    coords={
        "lon": (("y", "x"), lon_fy4a),
        "lat": (("y", "x"), lat_fy4a)
    }
)
ds.to_netcdf('ERA5_surface_vars_FY4A_2020010100.nc')
print("处理成功完成")


'''
裁剪
'''
# import xarray as xr
#
# # 读取地表数据（已裁剪）
# era5_surface = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/202001000_c.nc')
# sp = era5_surface['sp']  # 地表气压 (Pa)
# t2m = era5_surface['t2m']
# lon_era5 = era5_surface.longitude.values
# lat_era5 = era5_surface.latitude.values
# # 假设实际的时间坐标名称是 valid_time
# time = era5_surface.valid_time.values
#
# # 读取1000hPa分层比湿数据（需确保与地表数据维度一致）
# era5_q = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/specifichumidity_1000hPa_2020010100.nc')
# q = era5_q['q']
#
# # 提取 pressure_level = 1000hPa 这一层的数据
# q_1000hPa = q.sel(pressure_level=1000, drop=True)
#
# # 将 q_1000hPa 插值到地表数据的经纬度网格上
# q_interp = q_1000hPa.interp(latitude=lat_era5, longitude=lon_era5)
#
# # 创建新的 Dataset
# merged_dataset = xr.Dataset(
#     {
#         'sp': (('valid_time', 'latitude', 'longitude'), sp.data),
#         't2m': (('valid_time', 'latitude', 'longitude'), t2m.data),
#         'q': (('valid_time', 'latitude', 'longitude'), q_interp.data)
#     },
#     coords={
#         'valid_time': time,
#         'latitude': lat_era5,
#         'longitude': lon_era5
#     }
# )
#
# # 保存合并后的 Dataset 为 NetCDF 文件
# merged_dataset.to_netcdf('/mnt/datastore/liudddata/ERA5/20200101_Ttest/merged_data.nc')


